# -*- coding: utf-8 -*-  
'''
评价函数

@author: luoyi
Created on 2021年7月4日
'''
import tensorflow as tf

from models.crflinker.process import parse_y, padding_mask, viterbi


#    crf评价函数
class CRFMetrics(tf.metrics.Metric):
    def __init__(self,
                 name='crf_metrics',
                 crf_layer=None,
                 **kwargs):
        super(CRFMetrics, self).__init__(name=name, **kwargs)
        
        self._crf_layer = crf_layer
        
        self._acc = self.add_weight('ner_metrics', initializer='zero', dtype=tf.float64)
        pass
    
    def update_state(self, y_true, y_pred, sample_weight=None):
        '''计算词性标注预测正确的数量
            TP = et_to_et预测正确的数量
            TF = et_no预测正确的数量
            acc = TP + TF / P
        '''
        #    取标注预测
        x_pos = self._crf_layer.get_out()
        #    状态转移矩阵
        transfer = self._crf_layer.get_transfer()
        
        #    取标注值    Tensor(batch_size, max_handshake_size)，pad部分为-1
        y_id_pos, \
            _, _, _, \
            _, _, _ = parse_y(y_true)
        
        pad_mask = padding_mask(y_id_pos)
        #    vierbi算法返回的路径不包含首尾
        x_id_pos, _, _, _ = viterbi(x_pos, pad_mask, transfer)
        
        #    每个batch下标注数（包含[CLS][EOS]）
        P_B = tf.math.count_nonzero(y_id_pos >= 0, axis=-1)
        y_id_pos = tf.RaggedTensor.from_tensor(y_id_pos, padding=-1)
        #    y_id_pos相应的也去掉首尾[CLS][POS]
        y_id_pos = y_id_pos[:, 1:-1]
        
        eq = tf.math.equal(x_id_pos, y_id_pos)
        eq = tf.RaggedTensor.to_tensor(eq, default_value=False)
        TP_B = tf.math.count_nonzero(eq, axis=-1)
        
        #    计算准确率
        acc = tf.math.reduce_sum(TP_B) / tf.math.reduce_sum(P_B)
        self._acc.assign(acc)
        
        pass
    
    def reset_states(self):
        self._acc.assign(0.)
        pass
    def result(self):
        return self._acc
    
    pass



#    TPLinkerMetrics
class TPLinkerMetrics(tf.metrics.Metric):
    def __init__(self,
                 name='tplinker_metrics',
                 tplinker_layer=None,
                 **kwargs):
        super(TPLinkerMetrics, self).__init__(name=name, **kwargs)
        
        self._tplinker_layer = tplinker_layer
        
        self._acc = self.add_weight('re_metrics', initializer='zero', dtype=tf.float64)
        pass
    
    def update_state(self, y_true, y_pred, sample_weight=None):
        '''计算词性标注预测正确的数量
            TP = et_to_et预测正确的数量
            TF = et_no预测正确的数量
            acc = TP + TF / P
        '''
        #    取sht预测信息    Tensor(batch_size, rel_size, max_vec_len, 2, 3)
        x_sht = self._tplinker_layer.get_out()
        
        #    取标注值    Tensor(batch_size, max_handshake_size)，pad部分为-1
        _, \
            y_sh_1, y_sh_2, y_sh_no, \
            y_st_1, y_st_2, y_st_no = parse_y(y_true)
        
        #    计算正样本个数
        P_sh_1 = tf.math.count_nonzero(y_sh_1 >= 0)                 #    Tensor(batch_size, )
        P_sh_2 = tf.math.count_nonzero(y_sh_2 >= 0)                 #    Tensor(batch_size, )
        N_sh_no = tf.math.count_nonzero(y_sh_no >= 0)               #    Tensor(batch_size, )
        P_st_1 = tf.math.count_nonzero(y_st_1 >= 0)                 #    Tensor(batch_size, )
        P_st_2 = tf.math.count_nonzero(y_st_2 >= 0)                 #    Tensor(batch_size, )
        N_st_no = tf.math.count_nonzero(y_st_no >= 0)               #    Tensor(batch_size, )
        
        #    计算正样本预测结果正确数量
        TP_sh_1 = self.num(x_sht[:, :, :, 0, :], y_sh_1, 1)
        TP_sh_2 = self.num(x_sht[:, :, :, 0, :], y_sh_2, 2)
        TN_sh_no = self.num(x_sht[:, :, :, 0, :], y_sh_no, 0)
        TP_st_1 = self.num(x_sht[:, :, :, 1, :], y_st_1, 1)
        TP_st_2 = self.num(x_sht[:, :, :, 1, :], y_st_2, 2)
        TN_st_no = self.num(x_sht[:, :, :, 1, :], y_st_no, 0)
        
        #    计算准确率
        TP = TP_sh_1 + TP_sh_2 + TP_st_1 + TP_st_2
        TN = TN_sh_no + TN_st_no
        P = P_sh_1 + P_sh_2 + P_st_1 + P_st_2
        N = N_sh_no + N_st_no
        acc = (TP + TN) / (P + N)
        
        self._acc.assign(acc)
        
        pass
    
    def num(self, x, y, val):
        idx = tf.where(y[:, :, 0] >= 0)
        idx_B = idx[:, 0]
        idx_vec = tf.gather_nd(y, indices=idx)
        idx = tf.concat([idx_B[:, tf.newaxis], idx_vec], axis=-1)   #    Tensor(batch_size, 3)
        pred_sh_1 = tf.gather_nd(x, indices=idx)
        num = tf.math.count_nonzero(pred_sh_1 == val)
        return num
    
    def reset_states(self):
        self._acc.assign(0.)
        pass
    def result(self):
        return self._acc
    pass
